import os
import locale
import torch
from PIL import Image
from torchvision import transforms

# 设置系统区域环境，防止中文乱码
os.environ['LC_ALL'] = 'zh_CN.UTF-8'
locale.setlocale(locale.LC_ALL, 'zh_CN.UTF-8')

# 读取图片
def load_image(file_path):
    """加载并显示图片"""
    print(f"加载图片: {file_path}")
    image = Image.open(file_path).convert("RGB")  # 读取图片并转为RGB
    image.show()  # 显示图片
    print(f"图片尺寸为：{image.size}")  # 输出图片尺寸
    return image

def save_image(tensor, save_path):
    """保存张量类型的图片（需先转换为PIL图像）"""
    # 将张量转换为PIL图像（ToPILImage会自动处理维度和数据范围）
    pil_image = transforms.ToPILImage()(tensor)
    pil_image.save(save_path)
    print(f"图片已保存到: {save_path}")
    return save_path

# 定义图片预处理流程
def compose_transforms(five_crop_size=100):
    """定义并返回图片预处理流程"""
    # 单个图像的预处理（裁剪后对每个子图应用）
    single_transform = transforms.Compose([
        transforms.ColorJitter(brightness=0.5),  # 调整亮度
        transforms.Pad(padding=62,padding_mode="edge"),  # 边缘填充（示例值为10，可根据需要调整）
        transforms.ToTensor()  # 转为张量
    ])
    
    # 整体流程：先裁剪为5张图，再对每张图应用单个变换
    transform_pipeline = transforms.Compose([
        transforms.FiveCrop(size=five_crop_size),  # 生成5个裁剪图像（返回元组）
        # 对每个裁剪图像应用single_transform
        transforms.Lambda(lambda crops: torch.stack([single_transform(crop) for crop in crops]))
    ])
    
    print("图片预处理流程已定义")
    return transform_pipeline

def main(image_path, save_dir):
    # 创建保存目录（如果不存在）
    os.makedirs(save_dir, exist_ok=True)
    
    image = load_image(image_path)  # 加载图片
    transform_pipeline = compose_transforms()  # 获取预处理流程
    transformed_images = transform_pipeline(image)  # 应用预处理流程（返回形状为[5, 3, H, W]的张量）
    
    for i in range(transformed_images.size(0)):
        img_tensor = transformed_images[i]  # 取出第i个裁剪图像
        # 显示图像（需先转为PIL）
        transforms.ToPILImage()(img_tensor).show()
        print(f"第{i+1}张预处理后图片张量形状: {img_tensor.shape}")
        # 保存图像
        save_image(img_tensor, os.path.join(save_dir, f"transformed_image_{i+1}.png"))

if __name__ == "__main__":
    image_path = "../data/kaggle_cat_or_dog/training_data/cats/cat.2025.jpg"
    save_dir = "..\\tmp"
    main(image_path, save_dir)